from langchain_community.document_loaders import SQLDatabaseLoader
from langchain_community.utilities import SQLDatabase

from langchain_chroma import Chroma
from langchain_community.embeddings.huggingface import HuggingFaceBgeEmbeddings
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough
from langchain_core.prompts import ChatPromptTemplate
import os
from urllib.parse import quote_plus as urlquote
username = ""
password = ""
host = ""
port = ""
dbname = ""
# 连接字符串，mysql有mysql的方式，我这里是python的
pg_uri = f"postgresql+psycopg2://{username}:{urlquote(password)}@{host}:{port}/{dbname}"
db = SQLDatabase.from_uri(pg_uri)
# 加载sql 也可以通过csv等方式加载配置信息
loader = SQLDatabaseLoader(db=db, query="select * from table_name")
data = loader.load()
# 文本切割，如果文档不大可以不要切割，切割有时候也会导致出现问题
# text_splitter = RecursiveCharacterTextSplitter(
#     chunk_size=1000, chunk_overlap=200, add_start_index=True
# )
# all_splits = text_splitter.split_documents(data)
# 向量化存储
model_name = ""
model_kwargs = {}
encode_kwargs = {}
# embeding HuggingFace上有很多此处不公布我自己用的哪个
bgeEmbeddings = HuggingFaceBgeEmbeddings(
    model_name=model_name, model_kwargs=model_kwargs, encode_kwargs=encode_kwargs
)
# 此处跳过文本切割，因为切割后导致部分ID对应不起来的问题
db = Chroma.from_documents(data, bgeEmbeddings)
# 向量库检索
retriever = db.as_retriever(search_type="similarity", search_kwargs={'k': 5})
prompt = ChatPromptTemplate.from_template("""提出你的要求必要时可以提供样例
<context>
{context}
</context>
问题: {question}
""")
chat_zhipu = 自己实现
retriever_chain = (
        {"context": retriever, "question": RunnablePassthrough()}
        | prompt
        | chat_zhipu #此处随便挑一个langchain的llm实现即可，在这里我用的是智谱的chatglm4，因为我本地没法跑，问就是显卡不够
        | StrOutputParser()
)
if __name__ == '__main__':
    print(retriever.invoke("你的问题·······"))
